from openpyxl import Workbook
from openpyxl.drawing.image import Image
from openpyxl.utils import get_column_letter
from openpyxl.styles import Alignment
from pathlib import Path


class ExcelCreator:
    def __init__(self, directory="images/detect/"):
        self.directory = directory

    def create_excel_file(self, recognized_texts, filename):
        workbook = Workbook()
        sheet = workbook.active

        # 定义水平和垂直居中的对齐方式
        center_aligned_text = Alignment(horizontal='center', vertical='center')

        # 写入标题行并设置列宽
        titles = ["序号", "识别结果", "图片"]
        column_widths = [10, 30, 45]  # 调整每列的宽度
        for index, title in enumerate(titles, start=1):
            column_letter = get_column_letter(index)
            cell = sheet[column_letter + '1']
            cell.value = title
            cell.alignment = center_aligned_text
            sheet.column_dimensions[column_letter].width = column_widths[index - 1]

        # 写入数据行并尝试插入图片
        for row, item in enumerate(recognized_texts, start=1):
            row_number = row + 1
            # 设置序号单元格的值和居中对齐
            cell = sheet[f'A{row_number}']
            cell.value = row
            cell.alignment = center_aligned_text

            # 设置识别结果单元格的值和居中对齐
            cell = sheet[f'B{row_number}']
            cell.value = item["识别结果"]
            cell.alignment = center_aligned_text

            img_path = item["图片地址"]
            if Path(img_path).exists():
                img = Image(img_path)
                cell_ref = get_column_letter(3) + str(row_number)
                sheet.add_image(img, cell_ref)

            # 调整行高
            sheet.row_dimensions[row_number].height = 100

        # 保存文件
        Path(self.directory).mkdir(parents=True, exist_ok=True)
        filepath = Path(self.directory) / filename
        workbook.save(str(filepath))

        return str(filepath)


excel_creator = ExcelCreator()
